#!/usr/bin/env python
# coding=utf-8
# __author__ = 'Yunchao Ling'

from pymongo import MongoClient
import sys


# 连接MongoDB
def InitMongoDB(MONGODB_HOST, MONGODB_PORT, MONGODB_DB_NAME, MONGODB_COLLECTION_NAME):
    client = MongoClient(MONGODB_HOST, MONGODB_PORT)
    db = client[MONGODB_DB_NAME]
    collection = db[MONGODB_COLLECTION_NAME]
    return client, collection


# 关闭MongoDB
def CloseMongoDB(MongoDB_Connection):
    MongoDB_Connection.close()


def SearchChildren(collection, ht: set, max_rec: int, dms_dic: str, dms_id: str):
    result1 = collection.find_one({"dms_id": dms_id}, {"dms_id": 1, "dms_ids": 1})
    for item in result1["dms_ids"]:
        if len(ht) < max_rec:
            if item["db"] == dms_dic:
                if item["id"] not in ht:
                    ht.add(item["id"])
                    results2 = collection.find({"is_a": {"$elemMatch": {"db": dms_dic, "id": dms_id}}}, {"dms_id": 1})
                    for result2 in results2:
                        id = result2["dms_id"]
                        if len(ht) < max_rec:
                            SearchChildren(collection, ht, max_rec, dms_dic, id)


if __name__ == '__main__':
    db = "Pre_Data"
    coll = "SNOMEDCT_US"
    dms_dic = "SNOMEDCT_US"
    max_rec = 500
    ht = set()

    client, collection = InitMongoDB("10.188.188.22", 27017, db, coll)
    SearchChildren(collection, ht, max_rec, dms_dic, "6492006")
    CloseMongoDB(client)

    print(len(ht))
    for id in ht:
        print(dms_dic + ":" + id)
